import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.patches import Rectangle
import matplotlib.gridspec as gridspec
import sys

# 设置中文字体支持
plt.rcParams['font.sans-serif'] = ['SimHei', 'Arial Unicode MS', 'Microsoft YaHei', 'WenQuanYi Zen Hei', 'STHeiti']
plt.rcParams['axes.unicode_minus'] = False

# 设置参数
A = 1.0        # 振幅 (m)
omega = 4.0    # 角频率 (rad/s)
c = 2.0        # 波速 (m/s)
L = 10.0       # 弦的长度 (m)
T = 8.0        # 动画总时长 (s)
fps = 25       # 帧率

# 创建空间和时间网格
x = np.linspace(0, L, 1000)
num_frames = int(T * fps)
t_vals = np.linspace(0, T, num_frames)

# 预先计算波前位置
wavefront_positions = c * t_vals

# 创建图形
fig = plt.figure(figsize=(14, 10), facecolor='#f5f5f5')
fig.suptitle('半无限长弦受迫振动模拟', fontsize=20, fontweight='bold', color='#2c3e50')

# 使用GridSpec创建复杂的布局
gs = gridspec.GridSpec(3, 2, height_ratios=[3, 1, 1], width_ratios=[3, 1])
ax_main = plt.subplot(gs[0, 0])
ax_params = plt.subplot(gs[0, 1])
ax_wavefront = plt.subplot(gs[1, :])
ax_energy = plt.subplot(gs[2, :])

# 设置主图属性
ax_main.set_xlim(0, L)
ax_main.set_ylim(-A*1.8, A*1.8)
ax_main.set_xlabel('位置 x (m)', fontsize=12)
ax_main.set_ylabel('位移 u (m)', fontsize=12)
ax_main.grid(True, linestyle='--', alpha=0.7)
ax_main.axhline(0, color='k', lw=0.5)
ax_main.set_facecolor('#e8f4f8')

# 绘制端点振动器
ax_main.add_patch(Rectangle((-0.1, -1), 0.1, 2, color='#34495e', alpha=0.8))
ax_main.plot([-0.1, -0.1], [-1, 1], 'k-', lw=2)

# 参数区域设置
ax_params.axis('off')
params_text = ax_params.text(0.1, 0.5, '', fontsize=14, transform=ax_params.transAxes,
                             bbox=dict(boxstyle='round', facecolor='#ecf0f1', edgecolor='#bdc3c7', alpha=0.8))

# 波前传播图设置
ax_wavefront.set_title('波前传播', fontsize=14)
ax_wavefront.set_xlim(0, T)
ax_wavefront.set_ylim(0, L)
ax_wavefront.set_xlabel('时间 (s)', fontsize=12)
ax_wavefront.set_ylabel('位置 (m)', fontsize=12)
ax_wavefront.grid(True, linestyle='--', alpha=0.7)
ax_wavefront.set_facecolor('#f9ebea')

# 能量传播图设置
ax_energy.set_title('弦上能量分布', fontsize=14)
ax_energy.set_xlim(0, L)
ax_energy.set_ylim(0, 1.2)
ax_energy.set_xlabel('位置 x (m)', fontsize=12)
ax_energy.set_ylabel('相对能量', fontsize=12)
ax_energy.grid(True, linestyle='--', alpha=0.7)
ax_energy.set_facecolor('#eafaf1')

# 初始化波形
line, = ax_main.plot([], [], lw=3, color='#2980b9')
time_text = ax_main.text(0.02, 0.95, '', transform=ax_main.transAxes, fontsize=12, 
                        bbox=dict(facecolor='white', alpha=0.8))
wave_front = ax_main.scatter([], [], color='#e74c3c', s=80, zorder=5)
point_tracker = ax_main.scatter([], [], color='#9b59b6', s=80, zorder=5)

# 波前传播图初始化
wavefront_line, = ax_wavefront.plot([], [], 'r-', lw=2)
wavefront_point = ax_wavefront.scatter([], [], color='#e74c3c', s=80)

# 能量图初始化
energy_line, = ax_energy.plot([], [], 'g-', lw=2)

# 波动方程解
def wave_solution(x, t):
    """计算弦上各点的位移"""
    condition = t >= x / c
    return np.where(condition, A * np.sin(omega * (t - x / c)), 0)

# 计算弦上各点的动能和势能（相对值）
def wave_energy(x, t):
    """计算弦上各点的相对能量"""
    # 计算位移
    displacement = wave_solution(x, t)
    
    # 计算速度（通过时间导数）
    dt = 0.01
    if t <= dt:
        velocity = np.zeros_like(x)
    else:
        displacement_prev = wave_solution(x, t - dt)
        velocity = (displacement - displacement_prev) / dt
    
    # 计算应变（通过空间导数）
    dx = x[1] - x[0]
    strain = np.gradient(displacement, dx)
    
    # 动能和势能（相对值）
    kinetic_energy = 0.5 * velocity**2
    potential_energy = 0.5 * (c**2) * strain**2
    
    # 总能量（相对值）
    total_energy = kinetic_energy + potential_energy
    
    # 归一化处理
    max_energy = np.max(total_energy) if np.max(total_energy) > 0 else 1
    return total_energy / max_energy

# 动画更新函数
def update(frame_idx):
    """更新动画帧"""
    # 获取当前时间
    t = t_vals[frame_idx]
    
    # 计算弦的位移
    u = wave_solution(x, t)
    line.set_data(x, u)
    
    # 更新波前位置
    front_x = min(wavefront_positions[frame_idx], L)
    front_y = wave_solution(front_x, t) if front_x <= L else 0
    wave_front.set_offsets([[front_x, front_y]])
    
    # 更新跟踪点位置（选择几个特定位置的点）
    track_x = np.array([0, L/4, L/2, 3*L/4, L])
    track_y = wave_solution(track_x, t)
    point_tracker.set_offsets(np.c_[track_x, track_y])
    
    # 更新参数文本
    params_text.set_text(f'参数设置:\n\n振幅 A = {A} m\n角频率 ω = {omega} rad/s\n波速 c = {c} m/s\n波长 λ = {2*np.pi*c/omega:.2f} m\n周期 T = {2*np.pi/omega:.2f} s')
    
    # 更新波前传播图
    wavefront_line.set_data(t_vals[:frame_idx+1], wavefront_positions[:frame_idx+1])
    wavefront_point.set_offsets([[t, front_x]])
    
    # 更新能量图
    energy = wave_energy(x, t)
    energy_line.set_data(x, energy)
    
    # 更新时间文本
    time_text.set_text(f'时间: {t:.2f} s\n波前位置: {front_x:.2f} m')
    
    return (line, time_text, wave_front, point_tracker, params_text, 
            wavefront_line, wavefront_point, energy_line)

# 初始化函数 - 修复了散点图设置问题
def init():
    """动画初始化函数"""
    line.set_data([], [])
    time_text.set_text('')
    
    # 使用二维空数组初始化scatter的offsets
    wave_front.set_offsets(np.empty((0, 2)))  # 修复1: 使用正确形状的空数组
    point_tracker.set_offsets(np.empty((0, 2)))
    
    wavefront_line.set_data([], [])
    wavefront_point.set_offsets(np.empty((0, 2)))
    energy_line.set_data([], [])
    
    # 参数文本设为空
    params_text.set_text('')
    
    return (line, time_text, wave_front, point_tracker, params_text, 
            wavefront_line, wavefront_point, energy_line)

# 创建动画
print("正在创建动画...")
try:
    # 尝试使用更稳定的方式创建动画
    ani = FuncAnimation(
        fig, 
        update, 
        frames=num_frames, 
        init_func=init,
        blit=True, 
        interval=1000/fps,
        cache_frame_data=False  # 修复2: 禁用帧缓存以避免内存问题
    )
    
    # 保存为GIF
    print("正在保存动画为GIF文件...")
    try:
        ani.save('string_wave.gif', writer='pillow', fps=fps, dpi=100)
        print("动画已成功保存为 'string_wave.gif'")
    except Exception as e:
        print(f"保存动画时出错: {str(e)}")
        print("尝试使用较低的分辨率保存...")
        try:
            ani.save('string_wave_lowres.gif', writer='pillow', fps=fps, dpi=80)
            print("动画已成功保存为 'string_wave_lowres.gif' (较低分辨率)")
        except:
            print("无法保存动画，请检查依赖项和磁盘空间")
    
    plt.tight_layout(rect=[0, 0, 1, 0.95])
    plt.show()
    
except Exception as e:
    print(f"创建动画时发生错误: {str(e)}")
    print("尝试简化动画以解决问题...")
    
    # 简化版本 - 仅保留主视图
    fig, ax = plt.subplots(figsize=(10, 6))
    ax.set_xlim(0, L)
    ax.set_ylim(-A*1.5, A*1.5)
    ax.set_xlabel('位置 x (m)')
    ax.set_ylabel('位移 u (m)')
    ax.set_title('半无限长弦受迫振动模拟')
    ax.grid(True, linestyle='--', alpha=0.7)
    ax.axhline(0, color='k', lw=0.5)
    
    # 初始化波形
    line, = ax.plot([], [], lw=2, color='blue')
    time_text = ax.text(0.02, 0.95, '', transform=ax.transAxes)
    wave_front = ax.scatter([], [], color='red', s=50, zorder=5)
    
    # 简化更新函数
    def simple_update(frame_idx):
        t = t_vals[frame_idx]
        u = wave_solution(x, t)
        line.set_data(x, u)
        
        front_x = min(wavefront_positions[frame_idx], L)
        front_y = wave_solution(front_x, t) if front_x <= L else 0
        wave_front.set_offsets([[front_x, front_y]])
        
        time_text.set_text(f'时间: {t:.2f} s\n波前位置: {front_x:.2f} m')
        return line, time_text, wave_front
    
    # 简化初始化函数
    def simple_init():
        line.set_data([], [])
        time_text.set_text('')
        wave_front.set_offsets(np.empty((0, 2)))
        return line, time_text, wave_front
    
    # 创建简化动画
    ani = FuncAnimation(fig, simple_update, frames=num_frames, 
                        init_func=simple_init, blit=True, interval=1000/fps)
    
    # 保存简化动画
    try:
        ani.save('simple_string_wave.gif', writer='pillow', fps=fps, dpi=100)
        print("简化动画已成功保存为 'simple_string_wave.gif'")
    except Exception as e:
        print(f"保存简化动画时出错: {str(e)}")
    
    plt.show()